"""
股票资产模块 - 定义股票资产类及其特有功能
"""

from typing import Dict, List, Union, Any, Optional
import pandas as pd
from datetime import datetime, timedelta

from FDAS.data.assets.base import Asset
from FDAS.data.constants import ASSET_TYPE_STOCK
from FDAS.core.exceptions import DataFetchError

class Stock(Asset):
    """
    股票资产类
    
    提供股票特有的数据获取和分析功能
    """
    
    @property
    def asset_type(self) -> str:
        """获取资产类型"""
        return ASSET_TYPE_STOCK
    
    def _get_basic_info(self) -> Dict:
        """获取股票基本信息"""
        try:
            return self.provider.get_stock_basic_info(self.symbol)
        except Exception as e:
            return {
                "code": self.symbol,
                "name": "未知股票",
                "asset_type": self.asset_type,
                "error": str(e)
            }
    
    def get_history(self, 
                   start_date: Optional[str] = None, 
                   end_date: Optional[str] = None,
                   period: str = "daily", 
                   adjust: str = "qfq",
                   limit: Optional[int] = None) -> pd.DataFrame:
        """获取股票历史行情数据"""
        try:
            # 设置默认日期
            if not end_date:
                end_date = datetime.now().strftime("%Y%m%d")
            if not start_date:
                start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
            
            # 获取历史数据
            df = self.provider.get_stock_history(
                symbol=self.symbol,
                period=period,
                start_date=start_date,
                end_date=end_date,
                adjust=adjust
            )
            
            # 限制数据条数
            if limit and len(df) > limit:
                df = df.tail(limit)
                
            return df
        except Exception as e:
            raise DataFetchError(f"获取股票历史数据失败: {str(e)}")
    
    def get_valuation(self) -> pd.DataFrame:
        """获取股票估值数据"""
        try:
            # 获取原始估值数据
            df = self.provider.get_stock_valuation(self.symbol)
            
            # 如果数据为空，返回一个基本的DataFrame
            if df.empty:
                current_date = datetime.now().strftime('%Y-%m-%d')
                return pd.DataFrame({
                    '日期': [current_date],
                    '数据日期': [current_date],
                    'PE(TTM)': [None],
                    '市净率': [None],
                    '市销率': [None],
                    '市现率': [None]
                })
            
            # 创建一个新的DataFrame，只包含我们需要的列
            result_df = pd.DataFrame()
            
            # 处理日期列
            if 'trade_date' in df.columns:
                result_df['日期'] = df['trade_date']
            elif '日期' in df.columns:
                result_df['日期'] = df['日期']
            else:
                # 如果没有日期列，使用当前日期
                result_df['日期'] = datetime.now().strftime('%Y-%m-%d')
            
            # 处理PE(TTM)列
            if 'pe_ttm' in df.columns:
                result_df['PE(TTM)'] = df['pe_ttm']
            elif 'PE(TTM)' in df.columns:
                result_df['PE(TTM)'] = df['PE(TTM)']
            else:
                # 尝试其他可能的PE列名
                pe_cols = [col for col in df.columns if 'pe' in col.lower() or '市盈率' in col]
                if pe_cols:
                    result_df['PE(TTM)'] = df[pe_cols[0]]
                else:
                    result_df['PE(TTM)'] = None
            
            # 处理市净率列
            if 'pb' in df.columns:
                result_df['市净率'] = df['pb']
            elif '市净率' in df.columns:
                result_df['市净率'] = df['市净率']
            else:
                # 尝试其他可能的PB列名
                pb_cols = [col for col in df.columns if 'pb' in col.lower() or '市净' in col]
                if pb_cols:
                    result_df['市净率'] = df[pb_cols[0]]
                else:
                    result_df['市净率'] = None
            
            # 处理市销率列
            if 'ps' in df.columns:
                result_df['市销率'] = df['ps']
            elif '市销率' in df.columns:
                result_df['市销率'] = df['市销率']
            else:
                # 尝试其他可能的PS列名
                ps_cols = [col for col in df.columns if 'ps' in col.lower() or '市销' in col]
                if ps_cols:
                    result_df['市销率'] = df[ps_cols[0]]
                else:
                    result_df['市销率'] = None
            
            # 处理市现率列
            if 'pcf' in df.columns:
                result_df['市现率'] = df['pcf']
            elif '市现率' in df.columns:
                result_df['市现率'] = df['市现率']
            else:
                # 尝试其他可能的PCF列名
                pcf_cols = [col for col in df.columns if 'pcf' in col.lower() or '市现' in col]
                if pcf_cols:
                    result_df['市现率'] = df[pcf_cols[0]]
                else:
                    result_df['市现率'] = None
            
            # 添加数据日期列
            result_df['数据日期'] = result_df['日期']
            
            # 如果结果为空，返回一个基本的DataFrame
            if result_df.empty:
                current_date = datetime.now().strftime('%Y-%m-%d')
                return pd.DataFrame({
                    '日期': [current_date],
                    '数据日期': [current_date],
                    'PE(TTM)': [None],
                    '市净率': [None],
                    '市销率': [None],
                    '市现率': [None]
                })
            
            return result_df
        except Exception as e:
            raise DataFetchError(f"获取股票估值数据失败: {str(e)}")
    
    def get_financial(self, start_year: Optional[str] = None) -> Dict:
        """获取股票财务数据"""
        try:
            return self.provider.get_stock_financial(self.symbol)
        except Exception as e:
            raise DataFetchError(f"获取股票财务数据失败: {str(e)}")
    
    def get_fund_flow(self) -> Dict:
        """获取股票资金流向数据"""
        try:
            return self.provider.get_stock_fund_flow(self.symbol)
        except Exception as e:
            raise DataFetchError(f"获取股票资金流向数据失败: {str(e)}")
    
    def get_news(self, limit: int = 10) -> List[Dict]:
        """获取股票相关新闻"""
        try:
            return self.provider.get_stock_news(self.symbol)
        except Exception as e:
            raise DataFetchError(f"获取股票新闻数据失败: {str(e)}")
    
    def get_industry_info(self) -> Dict:
        """获取股票行业信息"""
        try:
            return self.provider.get_stock_industry_info(self.symbol)
        except Exception as e:
            raise DataFetchError(f"获取股票行业信息失败: {str(e)}")
    
    def get_shareholders(self) -> Dict:
        """获取股票股东信息"""
        try:
            return self.provider.get_stock_shareholders(self.symbol)
        except Exception as e:
            raise DataFetchError(f"获取股票股东信息失败: {str(e)}")
    
    def get_balance_sheet(self, limit: Optional[int] = None) -> pd.DataFrame:
        """
        获取资产负债表
        
        参数:
            limit: 限制返回的数据条数
            
        返回:
            包含资产负债表数据的DataFrame
        """
        try:
            # 获取财务数据
            financial_data = self.get_financial()
            
            # 提取资产负债表
            balance_sheet = financial_data.get('balance_sheet', pd.DataFrame())
            
            # 如果是空的DataFrame，创建一个基本的结构
            if balance_sheet.empty:
                balance_sheet = pd.DataFrame({
                    '报告期': [datetime.now().strftime('%Y-%m-%d')],
                    '资产总计': [None],
                    '负债合计': [None],
                    '所有者权益合计': [None]
                })
                return balance_sheet
            
            # 创建一个新的DataFrame，只包含我们需要的列
            result_df = pd.DataFrame()
            
            # 处理报告期列
            if '报告期' in balance_sheet.columns:
                result_df['报告期'] = balance_sheet['报告期']
            elif 'REPORT_DATE' in balance_sheet.columns:
                result_df['报告期'] = balance_sheet['REPORT_DATE']
            else:
                # 尝试查找包含"报告"或"日期"的列
                report_date_cols = [col for col in balance_sheet.columns if '报告' in col or '日期' in col]
                if report_date_cols:
                    result_df['报告期'] = balance_sheet[report_date_cols[0]]
                else:
                    # 如果没有找到合适的列，使用索引作为报告期
                    result_df['报告期'] = balance_sheet.index.astype(str)
            
            # 处理资产总计列
            if '资产总计' in balance_sheet.columns:
                result_df['资产总计'] = balance_sheet['资产总计']
            elif '*资产合计' in balance_sheet.columns:
                result_df['资产总计'] = balance_sheet['*资产合计']
            elif 'TOTAL_ASSETS' in balance_sheet.columns:
                result_df['资产总计'] = balance_sheet['TOTAL_ASSETS']
            else:
                # 尝试查找包含"资产"和"总计"或"合计"的列
                asset_cols = [col for col in balance_sheet.columns if '资产' in col and ('总计' in col or '合计' in col)]
                if asset_cols:
                    result_df['资产总计'] = balance_sheet[asset_cols[0]]
                else:
                    result_df['资产总计'] = None
            
            # 处理负债合计列
            if '负债合计' in balance_sheet.columns:
                result_df['负债合计'] = balance_sheet['负债合计']
            elif '*负债合计' in balance_sheet.columns:
                result_df['负债合计'] = balance_sheet['*负债合计']
            elif 'TOTAL_LIABILITIES' in balance_sheet.columns:
                result_df['负债合计'] = balance_sheet['TOTAL_LIABILITIES']
            else:
                # 尝试查找包含"负债"和"总计"或"合计"的列
                liability_cols = [col for col in balance_sheet.columns if '负债' in col and ('总计' in col or '合计' in col)]
                if liability_cols:
                    result_df['负债合计'] = balance_sheet[liability_cols[0]]
                else:
                    result_df['负债合计'] = None
            
            # 处理所有者权益合计列
            if '所有者权益合计' in balance_sheet.columns:
                result_df['所有者权益合计'] = balance_sheet['所有者权益合计']
            elif '*所有者权益（或股东权益）合计' in balance_sheet.columns:
                result_df['所有者权益合计'] = balance_sheet['*所有者权益（或股东权益）合计']
            elif 'TOTAL_EQUITY' in balance_sheet.columns:
                result_df['所有者权益合计'] = balance_sheet['TOTAL_EQUITY']
            else:
                # 尝试查找包含"权益"和"总计"或"合计"的列
                equity_cols = [col for col in balance_sheet.columns if '权益' in col and ('总计' in col or '合计' in col)]
                if equity_cols:
                    result_df['所有者权益合计'] = balance_sheet[equity_cols[0]]
                else:
                    result_df['所有者权益合计'] = None
            
            # 处理流动资产合计列
            if '流动资产合计' in balance_sheet.columns:
                result_df['流动资产合计'] = balance_sheet['流动资产合计']
            elif 'CURRENT_ASSETS' in balance_sheet.columns:
                result_df['流动资产合计'] = balance_sheet['CURRENT_ASSETS']
            else:
                # 尝试查找包含"流动资产"和"总计"或"合计"的列
                current_asset_cols = [col for col in balance_sheet.columns if '流动资产' in col and ('总计' in col or '合计' in col)]
                if current_asset_cols:
                    result_df['流动资产合计'] = balance_sheet[current_asset_cols[0]]
                else:
                    result_df['流动资产合计'] = None
            
            # 处理流动负债合计列
            if '流动负债合计' in balance_sheet.columns:
                result_df['流动负债合计'] = balance_sheet['流动负债合计']
            elif 'CURRENT_LIABILITIES' in balance_sheet.columns:
                result_df['流动负债合计'] = balance_sheet['CURRENT_LIABILITIES']
            else:
                # 尝试查找包含"流动负债"和"总计"或"合计"的列
                current_liability_cols = [col for col in balance_sheet.columns if '流动负债' in col and ('总计' in col or '合计' in col)]
                if current_liability_cols:
                    result_df['流动负债合计'] = balance_sheet[current_liability_cols[0]]
                else:
                    result_df['流动负债合计'] = None
            
            # 限制数据条数
            if limit and len(result_df) > limit:
                result_df = result_df.head(limit)
            
            return result_df
        except Exception as e:
            raise DataFetchError(f"获取资产负债表失败: {str(e)}")
    
    def get_income_statement(self, limit: Optional[int] = None) -> pd.DataFrame:
        """
        获取利润表
        
        参数:
            limit: 限制返回的数据条数
            
        返回:
            包含利润表数据的DataFrame
        """
        try:
            # 获取财务数据
            financial_data = self.get_financial()
            
            # 提取利润表
            income_statement = financial_data.get('income_statement', pd.DataFrame())
            
            # 如果是空的DataFrame，创建一个基本的结构
            if income_statement.empty:
                income_statement = pd.DataFrame({
                    '报告期': [datetime.now().strftime('%Y-%m-%d')],
                    '营业收入': [None],
                    '营业利润': [None],
                    '净利润': [None]
                })
                return income_statement
            
            # 创建一个新的DataFrame，只包含我们需要的列
            result_df = pd.DataFrame()
            
            # 处理报告期列
            if '报告期' in income_statement.columns:
                result_df['报告期'] = income_statement['报告期']
            elif 'REPORT_DATE' in income_statement.columns:
                result_df['报告期'] = income_statement['REPORT_DATE']
            else:
                # 尝试查找包含"报告"或"日期"的列
                report_date_cols = [col for col in income_statement.columns if '报告' in col or '日期' in col]
                if report_date_cols:
                    result_df['报告期'] = income_statement[report_date_cols[0]]
                else:
                    # 如果没有找到合适的列，使用索引作为报告期
                    result_df['报告期'] = income_statement.index.astype(str)
            
            # 处理营业收入列
            if '营业收入' in income_statement.columns:
                result_df['营业收入'] = income_statement['营业收入']
            elif '*营业总收入' in income_statement.columns:
                result_df['营业收入'] = income_statement['*营业总收入']
            elif 'TOTAL_REVENUE' in income_statement.columns:
                result_df['营业收入'] = income_statement['TOTAL_REVENUE']
            else:
                # 尝试查找包含"收入"的列
                revenue_cols = [col for col in income_statement.columns if '收入' in col]
                if revenue_cols:
                    result_df['营业收入'] = income_statement[revenue_cols[0]]
                else:
                    result_df['营业收入'] = None
            
            # 处理营业利润列
            if '营业利润' in income_statement.columns:
                result_df['营业利润'] = income_statement['营业利润']
            elif 'OPERATING_PROFIT' in income_statement.columns:
                result_df['营业利润'] = income_statement['OPERATING_PROFIT']
            else:
                # 尝试查找包含"营业"和"利润"的列
                profit_cols = [col for col in income_statement.columns if '营业' in col and '利润' in col]
                if profit_cols:
                    result_df['营业利润'] = income_statement[profit_cols[0]]
                else:
                    result_df['营业利润'] = None
            
            # 处理净利润列
            if '净利润' in income_statement.columns:
                result_df['净利润'] = income_statement['净利润']
            elif '*净利润' in income_statement.columns:
                result_df['净利润'] = income_statement['*净利润']
            elif 'NET_PROFIT' in income_statement.columns:
                result_df['净利润'] = income_statement['NET_PROFIT']
            else:
                # 尝试查找包含"净利润"的列
                net_profit_cols = [col for col in income_statement.columns if '净利润' in col]
                if net_profit_cols:
                    result_df['净利润'] = income_statement[net_profit_cols[0]]
                else:
                    result_df['净利润'] = None
            
            # 处理毛利列
            if '毛利' in income_statement.columns:
                result_df['毛利'] = income_statement['毛利']
            elif 'GROSS_PROFIT' in income_statement.columns:
                result_df['毛利'] = income_statement['GROSS_PROFIT']
            else:
                # 尝试查找包含"毛利"的列
                gross_profit_cols = [col for col in income_statement.columns if '毛利' in col and '率' not in col]
                if gross_profit_cols:
                    result_df['毛利'] = income_statement[gross_profit_cols[0]]
                else:
                    result_df['毛利'] = None
            
            # 处理毛利率列
            if '毛利率' in income_statement.columns:
                result_df['毛利率'] = income_statement['毛利率']
            elif 'GROSS_MARGIN' in income_statement.columns:
                result_df['毛利率'] = income_statement['GROSS_MARGIN']
            else:
                # 尝试查找包含"毛利率"的列
                gross_margin_cols = [col for col in income_statement.columns if '毛利率' in col]
                if gross_margin_cols:
                    result_df['毛利率'] = income_statement[gross_margin_cols[0]]
                else:
                    result_df['毛利率'] = None
            
            # 限制数据条数
            if limit and len(result_df) > limit:
                result_df = result_df.head(limit)
            
            return result_df
        except Exception as e:
            raise DataFetchError(f"获取利润表失败: {str(e)}")
    
    def get_cash_flow(self, limit: Optional[int] = None) -> pd.DataFrame:
        """
        获取现金流量表
        
        参数:
            limit: 限制返回的数据条数
            
        返回:
            包含现金流量表数据的DataFrame
        """
        try:
            # 获取财务数据
            financial_data = self.get_financial()
            
            # 提取现金流量表
            cash_flow = financial_data.get('cash_flow', pd.DataFrame())
            
            # 如果是空的DataFrame，创建一个基本的结构
            if cash_flow.empty:
                cash_flow = pd.DataFrame({
                    '报告期': [datetime.now().strftime('%Y-%m-%d')],
                    '经营活动产生的现金流量净额': [None],
                    '投资活动产生的现金流量净额': [None],
                    '筹资活动产生的现金流量净额': [None],
                    '现金及现金等价物净增加额': [None]
                })
            
            # 确保列名符合前端图表要求
            column_mapping = {
                'REPORT_DATE': '报告期',
                'NET_CASH_FLOW_OPERATING': '经营活动产生的现金流量净额',
                'NET_CASH_FLOW_INVESTING': '投资活动产生的现金流量净额',
                'NET_CASH_FLOW_FINANCING': '筹资活动产生的现金流量净额',
                'NET_CASH_FLOW': '现金及现金等价物净增加额'
            }
            
            # 重命名列
            for old_col, new_col in column_mapping.items():
                if old_col in cash_flow.columns and new_col not in cash_flow.columns:
                    cash_flow[new_col] = cash_flow[old_col]
            
            # 确保必要的列存在
            required_cols = ['报告期', '经营活动产生的现金流量净额', '投资活动产生的现金流量净额',
                           '筹资活动产生的现金流量净额', '现金及现金等价物净增加额']
            for col in required_cols:
                if col not in cash_flow.columns:
                    cash_flow[col] = None
            
            # 限制数据条数
            if limit and len(cash_flow) > limit:
                cash_flow = cash_flow.head(limit)
            
            return cash_flow
        except Exception as e:
            raise DataFetchError(f"获取现金流量表失败: {str(e)}")
    
    def get_financial_indicators(self, limit: Optional[int] = None) -> pd.DataFrame:
        """
        获取财务指标
        
        参数:
            limit: 限制返回的数据条数
            
        返回:
            包含财务指标数据的DataFrame
        """
        try:
            # 获取财务数据
            financial_data = self.get_financial()
            
            # 提取财务指标
            indicators = financial_data.get('indicators', pd.DataFrame())
            
            # 如果是空的DataFrame，创建一个基本的结构
            if indicators.empty:
                indicators = pd.DataFrame({
                    '报告期': [datetime.now().strftime('%Y-%m-%d')],
                    'ROE': [None],
                    'ROA': [None],
                    '资产负债率': [None],
                    '每股收益': [None],
                    '每股净资产': [None]
                })
            
            # 确保列名符合前端图表要求
            column_mapping = {
                'REPORT_DATE': '报告期',
                'RETURN_ON_EQUITY': 'ROE',
                'RETURN_ON_ASSETS': 'ROA',
                'DEBT_TO_ASSETS': '资产负债率',
                'EARNINGS_PER_SHARE': '每股收益',
                'NET_ASSETS_PER_SHARE': '每股净资产'
            }
            
            # 重命名列
            for old_col, new_col in column_mapping.items():
                if old_col in indicators.columns and new_col not in indicators.columns:
                    indicators[new_col] = indicators[old_col]
            
            # 确保必要的列存在
            required_cols = ['报告期', 'ROE', 'ROA', '资产负债率', '每股收益', '每股净资产']
            for col in required_cols:
                if col not in indicators.columns:
                    indicators[col] = None
            
            # 限制数据条数
            if limit and len(indicators) > limit:
                indicators = indicators.head(limit)
            
            return indicators
        except Exception as e:
            raise DataFetchError(f"获取财务指标失败: {str(e)}")
    
    def get_dividends(self) -> List[Dict]:
        """
        获取分红数据
        
        返回:
            包含分红数据的列表
        """
        try:
            # 这里应该调用provider的方法获取分红数据
            # 但由于当前没有实现，返回一个空列表
            return []
        except Exception as e:
            raise DataFetchError(f"获取分红数据失败: {str(e)}")
    
    def get_technical_indicators(self, limit: Optional[int] = None) -> Dict:
        """
        获取技术指标
        
        参数:
            limit: 限制返回的数据条数
            
        返回:
            包含技术指标的字典
        """
        try:
            # 获取历史数据
            history = self.get_history(limit=limit)
            
            # 如果历史数据为空，返回空字典
            if history.empty:
                return {}
            
            # 计算简单的技术指标
            result = {
                "ma5": [],
                "ma10": [],
                "ma20": [],
                "ma60": [],
                "dates": history['日期'].tolist() if '日期' in history.columns else []
            }
            
            # 确保有收盘价列
            if '收盘' in history.columns:
                # 计算移动平均线
                if len(history) >= 5:
                    history['ma5'] = history['收盘'].rolling(window=5).mean()
                    result["ma5"] = history['ma5'].tolist()
                
                if len(history) >= 10:
                    history['ma10'] = history['收盘'].rolling(window=10).mean()
                    result["ma10"] = history['ma10'].tolist()
                
                if len(history) >= 20:
                    history['ma20'] = history['收盘'].rolling(window=20).mean()
                    result["ma20"] = history['ma20'].tolist()
                
                if len(history) >= 60:
                    history['ma60'] = history['收盘'].rolling(window=60).mean()
                    result["ma60"] = history['ma60'].tolist()
            
            return result
        except Exception as e:
            raise DataFetchError(f"获取技术指标失败: {str(e)}")